import argparse

import yaml
from ultralytics import YOLO
from user_det_sg2yolo import sengo2yolo


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--root_path', type=str, default=r'D:\data\231207huoni\trainV8Seg_cable', help='root_path')
    par = parser.parse_args()
    root_path = par.root_path

    # data
    # convert data todo
    data_yaml_path = f'{root_path}/format_data/data.yaml'

    # train
    project = f'{root_path}/models'
    # weights = f'{root_path}/models/train6/weights/last.pt'
    model = YOLO('yolov8n-seg.yaml').load('yolov8n-seg.pt')   # build from YAML and transfer weights
    # model = YOLO()
    # results = model.train(data=data_yaml_path, epochs=100, imgsz=640, device=1,
    #                       batch=-1, workers=0, project=project)

    results = model.train(data=data_yaml_path,  epochs=100,
                          imgsz=640, device=0, resume=True,
                          batch=-1, workers=0, project=project,
                          name='v8sSeg_640_cable', overlap_mask=False)